import os
import sys
import json
sys.path.append(r'src')
from data.api import iter_dataset

# simple script to re-arrange CFG nodes by addr in json
# (Modify in place)

dataset_dir = r'preprocessed_data/cfg_transferred/CWE121_Stack_Based_Buffer_Overflow'
all_files = iter_dataset(dataset_dir)
for files,label in all_files:
    cfg = files[1]
    d = {}  
    ls = []  
    f = open(cfg, 'r')
    con = json.load(f)
    for n in con["nodes"]:
        d[n["id"]] = n
    d = dict(sorted(d.items(),key=lambda x:x[0]))
    for id in d.keys():
        ls.append(d[id])
    f.close()
    f = open(cfg, 'w')
    json.dump({"nodes":ls, "edges":con["edges"]}, f)